from sklearn.datasets import fetch_mldata
import matplotlib.pyplot as plt
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import precision_score
import matplotlib
import numpy as np
# 获取数据
mnist = fetch_mldata('mnist-original',data_home='./')
x,y = mnist['data'],mnist['target']
# 划分训练集测试集
x_train,x_test,y_train,y_test = x[:60000,:],x[60000:,:],y[:60000],y[60000:]
shuffer_index = np.random.permutation(60000)
x_train,y_train=x_train[shuffer_index],y_train[shuffer_index]
x_train,y_train = x_train[:60000],y_train[:60000]
x_train_array = []
train_size = x_train.shape[0]
# 遍历训练集
for i in range(train_size):
    x_train_array.append(x_train[i].reshape(28,28))

def data_en(datas,direct='u'):
    size = len(datas)
    ret_en = np.zeros((size,784))
    # print(size,ret_en)
    if direct == 'u':
        for i in range(size):
            # 定义空列表trans_data,先加第二行到最后一行数据，再加第一行数据
            trans_data = np.append(datas[i][1:,:],datas[i][0:1,:],axis=0)
            # print(trans_data,trans_data.shape)
            # plt.imshow(trans_data, cmap=matplotlib.cm.binary, interpolation='nearest')
            # plt.show()
            # 将数据转换成一行
            ret_en[i] = trans_data.reshape(1,-1)
    elif direct == 'd':
        for i in range(size):
            trans_data = np.append(datas[i][-1:,:],datas[i][:-1,:],axis=0)
            # plt.imshow(trans_data, cmap=matplotlib.cm.binary, interpolation='nearest')
            # plt.show()
            # print(trans_data,trans_data.shape)
            ret_en[i] = trans_data.reshape(1,-1)
            # print(ret_en[i].shape)
    elif direct == 'l':
        for i in range(size):
            trans_data = np.append(datas[i][:,1:],datas[i][:,0:1],axis=1)
            # plt.imshow(trans_data, cmap=matplotlib.cm.binary, interpolation='nearest')
            # plt.show()
            # print(trans_data,trans_data.shape)
            ret_en[i] = trans_data.reshape(1,-1)
            # print(ret_en[i].shape)
    elif direct == 'r':
        for i in range(size):
            trans_data = np.append(datas[i][:,-1:],datas[i][:,:-1],axis=1)
            # plt.imshow(trans_data, cmap=matplotlib.cm.binary, interpolation='nearest')
            # plt.show()
            # print(trans_data,trans_data.shape)

    return ret_en

x_train_u = data_en(x_train_array)
x_train_d = data_en(x_train_array,'d')
x_train_l = data_en(x_train_array,'l')
x_train_r = data_en(x_train_array,'r')

x_train_A = np.concatenate((x_train,x_train_d,x_train_l,x_train_u,x_train_r),axis=0)
y_train_A = np.concatenate((y_train,y_train,y_train,y_train,y_train),axis=0)

# 测试前打分
knnf = KNeighborsClassifier()
knnf.fit(x_train,y_train)
y_pred = knnf.predict(x_test)
ps = precision_score(y_test,y_pred,average=None)
print('\nps=',ps,np.average(ps))
# 测试后打分
knnf_A = KNeighborsClassifier()
knnf_A.fit(x_train_A,y_train_A)
y_pred_A = knnf.predict(x_test)
ps_A = precision_score(y_test,y_pred_A,average=None)
print('\nps_A=',ps_A,np.average(ps_A))